package com.dify.knowledge.util;

import lombok.extern.slf4j.Slf4j;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.apache.poi.hwpf.HWPFDocument;
import org.apache.poi.hwpf.extractor.WordExtractor;
import org.apache.poi.ss.usermodel.*;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import org.apache.poi.xwpf.usermodel.XWPFParagraph;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;

/**
 * 文档文本提取工具类
 * 支持PDF、Word、Excel、TXT格式的文本提取
 */
@Slf4j
public class DocumentTextExtractor {
    
    /**
     * 根据文件类型提取文本内容
     * @param file 文件对象
     * @param fileType 文件类型：PDF, WORD, EXCEL, TXT
     * @return 提取的文本内容
     * @throws IOException 文件读取异常
     */
    public static String extractText(File file, String fileType) throws IOException {
        if (file == null || !file.exists()) {
            throw new IllegalArgumentException("文件不存在");
        }
        
        return switch (fileType.toUpperCase()) {
            case "PDF" -> extractFromPdf(file);
            case "WORD" -> extractFromWord(file);
            case "EXCEL" -> extractFromExcel(file);
            case "TXT" -> extractFromTxt(file);
            default -> throw new IllegalArgumentException("不支持的文件类型: " + fileType);
        };
    }
    
    /**
     * 从PDF文件提取文本
     * @param file PDF文件
     * @return 提取的文本内容
     * @throws IOException 文件读取异常
     */
    public static String extractFromPdf(File file) throws IOException {
        log.info("开始提取PDF文本: {}", file.getName());
        
        try (PDDocument document = PDDocument.load(file)) {
            PDFTextStripper stripper = new PDFTextStripper();
            String text = stripper.getText(document);
            log.info("PDF文本提取完成，共{}页，文本长度: {}", document.getNumberOfPages(), text.length());
            return text;
        }
    }
    
    /**
     * 从Word文件提取文本
     * 支持.doc和.docx格式
     * @param file Word文件
     * @return 提取的文本内容
     * @throws IOException 文件读取异常
     */
    public static String extractFromWord(File file) throws IOException {
        log.info("开始提取Word文本: {}", file.getName());
        
        String fileName = file.getName().toLowerCase();
        String text;
        
        if (fileName.endsWith(".docx")) {
            // 处理.docx格式
            try (FileInputStream fis = new FileInputStream(file);
                 XWPFDocument document = new XWPFDocument(fis)) {
                StringBuilder sb = new StringBuilder();
                for (XWPFParagraph paragraph : document.getParagraphs()) {
                    sb.append(paragraph.getText()).append("\n");
                }
                text = sb.toString();
            }
        } else if (fileName.endsWith(".doc")) {
            // 处理.doc格式
            try (FileInputStream fis = new FileInputStream(file);
                 HWPFDocument document = new HWPFDocument(fis);
                 WordExtractor extractor = new WordExtractor(document)) {
                text = extractor.getText();
            }
        } else {
            throw new IllegalArgumentException("不支持的Word文件格式: " + fileName);
        }
        
        log.info("Word文本提取完成，文本长度: {}", text.length());
        return text;
    }
    
    /**
     * 从Excel文件提取文本
     * 支持.xls和.xlsx格式
     * @param file Excel文件
     * @return 提取的文本内容
     * @throws IOException 文件读取异常
     */
    public static String extractFromExcel(File file) throws IOException {
        log.info("开始提取Excel文本: {}", file.getName());
        
        StringBuilder sb = new StringBuilder();
        
        try (FileInputStream fis = new FileInputStream(file);
             Workbook workbook = WorkbookFactory.create(fis)) {
            
            int numberOfSheets = workbook.getNumberOfSheets();
            log.info("Excel文件共有{}个工作表", numberOfSheets);
            
            for (int i = 0; i < numberOfSheets; i++) {
                Sheet sheet = workbook.getSheetAt(i);
                sb.append("工作表: ").append(sheet.getSheetName()).append("\n");
                
                for (Row row : sheet) {
                    for (Cell cell : row) {
                        String cellValue = getCellValueAsString(cell);
                        if (cellValue != null && !cellValue.trim().isEmpty()) {
                            sb.append(cellValue).append("\t");
                        }
                    }
                    sb.append("\n");
                }
                sb.append("\n");
            }
        }
        
        String text = sb.toString();
        log.info("Excel文本提取完成，文本长度: {}", text.length());
        return text;
    }
    
    /**
     * 从TXT文件提取文本
     * 支持UTF-8编码的纯文本文件
     * @param file TXT文件
     * @return 提取的文本内容
     * @throws IOException 文件读取异常
     */
    public static String extractFromTxt(File file) throws IOException {
        log.info("开始提取TXT文本: {}", file.getName());
        
        try {
            String text = Files.readString(file.toPath(), StandardCharsets.UTF_8);
            log.info("TXT文本提取完成，文本长度: {}", text.length());
            return text;
        } catch (IOException e) {
            // 如果UTF-8读取失败，尝试使用系统默认编码
            log.warn("UTF-8编码读取失败，尝试使用系统默认编码: {}", e.getMessage());
            try {
                String text = Files.readString(file.toPath());
                log.info("TXT文本提取完成（系统默认编码），文本长度: {}", text.length());
                return text;
            } catch (IOException ex) {
                log.error("TXT文件读取失败: {}", ex.getMessage());
                throw ex;
            }
        }
    }
    
    /**
     * 获取单元格的字符串值
     * @param cell 单元格对象
     * @return 单元格的字符串值
     */
    private static String getCellValueAsString(Cell cell) {
        if (cell == null) {
            return "";
        }
        
        return switch (cell.getCellType()) {
            case STRING -> cell.getStringCellValue();
            case NUMERIC -> {
                if (DateUtil.isCellDateFormatted(cell)) {
                    yield cell.getDateCellValue().toString();
                } else {
                    yield String.valueOf(cell.getNumericCellValue());
                }
            }
            case BOOLEAN -> String.valueOf(cell.getBooleanCellValue());
            case FORMULA -> cell.getCellFormula();
            case BLANK -> "";
            default -> "";
        };
    }
}
